
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
class Ic_record():
    def __init__(self, args, codes) -> None:
        self.args = args
        self.codes = codes
                 
    def ic_record(self, pred_df, true_df, times, name='test'):  
        ic_list, rmse_list = record_calc(pred_df.values, true_df.values)
        print( f'当前样本外IC_mean：{np.nanmean(ic_list)}, IC_std：{np.nanstd(ic_list)},rmse：{np.nanmean(rmse_list)}')
        if self.args.predict_save:
            pred_df.to_csv(f'{self.args.result_root}/{name}_pred.csv')
            pred_df.to_pickle(f'{self.args.result_root}/{name}_pred.pkl.gzip')
            true_df.to_csv(f'{self.args.result_root}/{name}_true.csv')
            true_df.to_pickle(f'{self.args.result_root}/{name}_true.pkl.gzip')     
        ic_df = pd.DataFrame(index=self.codes)
        ic_df['ic'] = ic_list
        ic_df['rmse'] = rmse_list
        ic_values_id = ic_df.values 
        mean_id = np.nanmean(ic_values_id, axis = 0)
        std_id = np.nanstd(ic_values_id, axis = 0)
        ic_df.loc['mean'] = mean_id
        ic_df.loc['std'] = std_id
        ic_df.to_csv(f'{self.args.result_root}/{name}_ic.csv')       
        plot_lint_ml(array1=pred_df.iloc[-1000:,0].values, array2=true_df.iloc[-1000:,0].values, times=times,  name=name, root=self.args.result_root)
 
 
def record_calc(pred,true):
    ic_list = []
    for i in range(pred.shape[1]):
        cor = calc_corr(pred[:, i], true[:, i])
        ic_list.append(cor)
    rmse_list = []
    for i in range(pred.shape[1]):
        rmse = calc_rmse(pred[:, i], true[:, i])
        rmse_list.append(rmse)
    return ic_list, rmse_list

def calc_rmse(array_a, array_b):
    return np.sqrt(np.nansum(np.square(array_a-array_b)))/len(array_a)

def calc_corr(array_a, array_b):
    if np.nanstd(array_b.tolist())==0:
        print('label无波动，相关系数计算失败')
    if np.nanstd(array_a.tolist())==0:
        print('预测结果无波动，相关系数计算失败')
    return (np.nanmean((array_a*array_b).tolist())-np.nanmean(array_a.tolist())*np.nanmean(array_b.tolist()))/(np.nanstd(array_a.tolist())*np.nanstd(array_b.tolist()))


def plot_lint_ml(array1, array2, times, name, root):
    plt.figure(figsize=(30, 10))
    plt.plot(array1,  alpha=0.5, label='pred')
    plt.plot(array2, alpha=0.5, label='true')
    plt.plot(np.ones(len(array1))*0.01, label = '0.01' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*0.02, label = '0.02' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*0.05, label = '0.05' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*-0.01, label = '-0.01' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*-0.02, label = '-0.02' , alpha=0.5, linestyle=':')
    plt.plot(np.ones(len(array1))*-0.05, label = '-0.05' , alpha=0.5, linestyle=':')
    plt.legend()
    save_root = f'{root}/{name}_result.png'
    plt.savefig(save_root)
    plt.close()
